import numpy as np
import matplotlib.pyplot as plt
from scipy.fft import fft, fftfreq, fftshift
from scipy import signal

F = 5.
T = 20.
N = 10001
DT = T/N
t = np.linspace(0.,T,N)
t[0:5]

y_sq =0.5+0.5*signal.square(2*np.pi*F*t)
plt.subplot(2, 1, 1)
plt.plot(t,y_sq)
plt.grid(True)
plt.xlabel("time[s]",position=(0.95,1))
plt.ylabel("signal[Volt]",position=(1,0.8))
ax = plt.gca()
ax.set_xlim(0,1)

y_sq_fft = fftshift(fft(y_sq))
x_fft = fftshift(fftfreq(N,DT))

plt.subplot(2, 1, 2)
plt.semilogy(x_fft,2./N*np.abs(y_sq_fft))
plt.grid(True)
ax = plt.gca()
ax.set_xlim(0,100)
ax.set_ylim(1e-4,1)
plt.xlabel("Frequency[Hz]",position=(0.9,1))
plt.ylabel("Signal[Volt]",position=(1,0.9))

plt.subplots_adjust(left=0.1, right=0.9, 
                    top=0.9, bottom=0.1, 
                    wspace=0.4, hspace=0.4)
plt.show()

